#!/usr/bin/python
# -*- coding:utf-8 -*-
# @Time  : 6/1/20
# @Author: yanlh
# @usage : docker容器插件调用工具
import os
import sys

sys.path.append('..')
from common_utils import CBBCommonUtils

"""
module: common parameters
"""
rein_info = ['Please check and config the unsafe item manually.', 'Reinforce Complete.']
rollback_info = ["This function don't need to rollback.", 'Rollback Complete.']
backup_suffix = 'BACKUP.AQJSB'
"""
module: daemon file config
func: 
    1. validate file/catalog exist
    2. validate file owner/privilege
    3. backup file
    4. rollback file
"""


def validate_fragment_file_exist(func_args, target=None):
    filename = func_args.get('filename', None)
    if target:
        filename = target.get('path', None)
    cmd = 'systemctl show -p FragmentPath {0}'.format(filename)
    result, output = CBBCommonUtils.cbb_run_cmd(func_args['ip'], cmd, username=func_args['username'],
                                                passwd=func_args['passwd'])
    if result:
        des = "CMD process ERROR. Error info: {}".format(result)
        return 'error', des
    for item in output:
        if 'FragmentPath' in item:
            value = list(filter(None, item.split('=')))
            if len(value) > 1 and value[1] and value[1].strip():
                return 'success', value[1].strip()
        continue
    des = cmd + " is NOT EXIST."
    return 'info', des


def validate_catalog_file_exist(func_args):
    cmd = 'find {}'.format(func_args['path'])
    result, output = CBBCommonUtils.cbb_run_cmd(func_args['ip'], cmd, username=func_args['username'],
                                                passwd=func_args['passwd'])
    if result:
        des = "CMD process ERROR. Error info: {}".format(result)
        return 'error', des
    if output and len(output):
        return 'success', None
    else:
        des = func_args['path'] + " is NOT EXIST."
        return 'info', des


def valid_file_exist(func_args):
    if func_args.get('file', None) == 'fragment':
        valid, des = validate_fragment_file_exist(func_args)
    else:
        valid, des = validate_catalog_file_exist(func_args)
    return valid, des


def validate_file_owner(path, func_args, is_get=False):
    """
    校验文件的所有者是否为owner
    """
    cmd = "stat -c %U:%G {0}".format(path, func_args['owner'])
    result, output = CBBCommonUtils.cbb_run_cmd(func_args['ip'], cmd, username=func_args['username'],
                                                passwd=func_args['passwd'])
    if result:
        des = "CMD process ERROR. Error info: {}".format(result)
        info = {'info': des, 'type': 'error'}
        return 'error', info
    value = ''.join(output).strip() if output else None
    if is_get:
        return 'success', value
    if func_args['owner'] == value:
        des = 'Safe. The owner of file {} is safe'.format(path)
        info = {'info': des, 'type': 'info'}
        return 'success', info
    des = 'Not Safe. File {} is not safe, the owner should be {}'.format(path, func_args['owner'])
    info = {'info': des, 'type': 'warning'}
    return 'fail', info


def validate_docker_file_privilege(path, func_args, is_get=False):
    """
    校验文件的权限
    """
    cmd = "stat -c %a {0}".format(path)
    result, output = CBBCommonUtils.cbb_run_cmd(func_args['ip'], cmd, username=func_args['username'],
                                                passwd=func_args['passwd'])
    if result:
        des = "CMD process ERROR. Error info: {}".format(result)
        info = {'info': des, 'type': 'error'}
        return 'error', info
    value = ''.join(output).strip() if output else None
    if is_get:
        return 'success', value
    if int(func_args['privilege']) >= int(value):
        des = 'Safe. The privilege of file {} is safe'.format(path)
        info = {'info': des, 'type': 'info'}
        return 'success', info
    des = 'Not Safe. File {} is not safe, the privilege should be {}'.format(path, func_args['privilege'])
    info = {'info': des, 'type': 'warning'}
    return 'fail', info


def modify_file_owner(path, func_args):
    cmd = 'chown {} {}'.format(func_args['owner'], path)
    result, output = CBBCommonUtils.cbb_run_cmd(func_args['ip'], cmd, username=func_args['username'],
                                                passwd=func_args['passwd'])
    return not result


def modify_file_privilege(path, func_args):
    cmd = 'chmod {} {}'.format(func_args['privilege'], path)
    result, output = CBBCommonUtils.cbb_run_cmd(func_args['ip'], cmd, username=func_args['username'],
                                                passwd=func_args['passwd'])
    return not result


def check_backup_file(func_args, target, backup=None):
    if backup:
        path, filename = os.path.split(backup)
        cmd = "ls {} |grep {}".format(path, filename)
    else:
        path, filename = os.path.split(target)
        cmd = "ls {} | grep {}.{}.*".format(path, filename, backup_suffix)
    result, output = CBBCommonUtils.cbb_run_cmd(func_args['ip'], cmd, username=func_args['username'],
                                                passwd=func_args['passwd'])
    if result:
        des = 'CMD process ERROR. Error info: {}'.format(result)
        info = {'info': des, 'type': 'error'}
        return False, info
    if output and len(output):
        backup_name = output[len(output) - 1].replace("\n", "")
        value = "{}/{}".format(path, backup_name)
        return True, value
    else:
        return True, None


def write_backup_file(func_args, target, priv_info):
    content = '{};{}'.format(priv_info[0], priv_info[1])
    cmd = "echo '{}' > {}".format(content, target)
    result, output = CBBCommonUtils.cbb_run_cmd(func_args['ip'], cmd, username=func_args['username'],
                                                passwd=func_args['passwd'])
    if result:
        des = 'Write info failed. Error info: {}'.format(result)
        info = {'info': des, 'type': 'warning'}
        return False, info
    else:
        des = {'info': 'Write info success.', 'type': 'info'}
        return True, des


def create_backup_file(func_args, target):
    cmd = 'touch {}'.format(target)
    result, output = CBBCommonUtils.cbb_run_cmd(func_args['ip'], cmd, username=func_args['username'],
                                                passwd=func_args['passwd'])
    if result:
        des = 'CMD process ERROR. Error info: {}'.format(result)
        info = {'info': des, 'type': 'error'}
        return False, info
    return True, None


def record_owner_priv_info(func_args, target):
    fmark, priv_info = get_origin_owner_priv(func_args, target)
    if not fmark:
        return False, priv_info
    target_backup = "{}.{}.`date +%Y-%m-%d`".format(target, backup_suffix)
    create_flag, create_result = create_backup_file(func_args, target_backup)
    if not create_flag:
        return False, [create_result]
    check_flag, check_result = check_backup_file(func_args, target, target_backup)
    if not check_flag:
        return False, [check_result]
    if not check_result:
        des = 'File {} create failed.'.format(target)
        info = {'info': des, 'type': 'warning'}
        return False, [info]
    write_falg, write_result = write_backup_file(func_args, target_backup, priv_info)
    if not write_falg:
        delete_backup_file(func_args, target_backup)
    return write_falg, [write_result]


def get_backup_file(func_args, target):
    flag, result = check_backup_file(func_args, target)
    if not flag:
        return result, 1
    if not result:
        info = {'info': 'Backup file NOT FOUND.', 'type': 'warning'}
        return info, -1
    return result, 0


def get_origin_owner_priv(func_args, file):
    owner_flag, owner_value = validate_file_owner(file, func_args, True)
    priv_flag, priv_value = validate_docker_file_privilege(file, func_args, True)
    if owner_flag != 'success' or priv_flag != 'success':
        return False, [owner_value, priv_value]
    else:
        return True, [owner_value, priv_value]


def get_backup_owner_priv(func_args, target):
    check_flag, check_reuslt = check_backup_file(func_args, target)
    if not check_flag and not check_reuslt:
        return False, check_reuslt
    cmd = 'cat {}'.format(check_reuslt)
    result, output = CBBCommonUtils.cbb_run_cmd(func_args['ip'], cmd, username=func_args['username'],
                                                passwd=func_args['passwd'])
    if result:
        des = "CMD process ERROR. Error info: {}".format(result)
        info = {'info': des, 'type': 'error'}
        return False, info
    value = ''.join(output)
    if not value or not value.strip():
        des = 'Nothing in Backup file'
        return False, {'info': des, 'type': 'warning'}
    value = value.strip().split(';')
    if len(value) < 2:
        des = 'Owner and privilege in Backup is wrong.'
        return False, {'info': des, 'type': 'warning'}
    owner = value[0].strip() if value[0] else None
    privilege = value[1].strip() if value[1] else None
    return True, [owner, privilege]


def modify_target_privilege(func_args, target, priv_info):
    func_args['owner'] = priv_info[0]
    func_args['privilege'] = priv_info[1]
    owner = modify_file_owner(target, func_args)
    priv = modify_file_privilege(target, func_args)
    if not owner or not priv:
        des = {'info': "Rollback FAILED, please retry.", 'type': 'error'}
        return des, 1
    des = {'info': 'Rollback success.', 'type': 'info'}
    return des, 0


def delete_backup_file(func_args, target):
    if not os.path.isfile(target):
        des = 'File {} is not a file, cannot be deleted.'.format(target)
        return {'info': des, 'type': 'warning'}, 1
    cmd = "rm -rf {}".format(target)
    result, output = CBBCommonUtils.cbb_run_cmd(func_args['ip'], cmd, username=func_args['username'],
                                                passwd=func_args['passwd'])
    if not result:
        des = "Backup file {} delete Success.".format(target)
        info = {'info': des, 'type': 'info'}
        return info, 0
    else:
        des = 'Backup file {} delete Failed.'.format(target)
        info = {'info': des, 'type': 'warning'}
        return info, 1


def rollback_file(func_args):
    # 查找备份文件，此处的查找支持多备份文件的查找（结果会按时间顺序排序），并默认回滚最新一次备份的内容，可以连续多次回滚
    valid, target = valid_file_exist(func_args)
    if valid != 'success':
        record_list, scan_result = scan_fail_file_exist(valid, target)
        return record_list, scan_result.get('error_count')
    if not target:
        target = func_args.get('path', None)
    fmark, priv_info = get_backup_owner_priv(func_args, target)
    if not fmark:
        return [priv_info], 1
    info, error = get_backup_file(func_args, target)
    if error != 0:
        return [info], error
    # 回滚文件
    rollback, error_count = modify_target_privilege(func_args, target, priv_info)
    if error_count:
        return [rollback], error_count
    delete, error_count = delete_backup_file(func_args, info)
    return [rollback, delete], error_count


def scan_fail_file_exist(valid, des):
    record_list = []
    error_count = 0
    record_list.append({'info': des, 'type': 'warning'})
    if valid == 'error':
        error_count += 1
    record_list.append({'info': 'Scan Complete.', 'type': 'info'})
    scan_result = {'flag': False, 'error_count': error_count}
    return record_list, scan_result


def scan_success_file_exist(path, func_args):
    record_list = []
    error_count = 0
    if not path:
        path = func_args.get('path', None)
    owner, info_owner = validate_file_owner(path, func_args)
    priv, info_priv = validate_docker_file_privilege(path, func_args)
    record_list.append(info_owner)
    record_list.append(info_priv)
    is_reinforce = False
    if owner != 'success' or priv != 'success':
        error_count += 1
    if owner == 'fail' or priv == 'fail':
        is_reinforce = True
    record_list.append({'info': 'Scan Complete.', 'type': 'info'})
    scan_result = {'flag': is_reinforce, 'error_count': error_count, 'owner': owner, 'priv': priv, 'path': path}
    return record_list, scan_result


def scan_file_config(func_args):
    valid, des = valid_file_exist(func_args)
    if valid != 'success':
        record_list, scan_result = scan_fail_file_exist(valid, des)
        return record_list, scan_result
    record_list, scan_result = scan_success_file_exist(des, func_args)
    return record_list, scan_result


def reinforce_file_config(func_args, scan_result):
    record_list = []
    error_count = 0
    path = scan_result.get('path', None)
    if not path:
        path = func_args.get('path', None)
    backup, backup_info = record_owner_priv_info(func_args, path)
    record_list.extend(backup_info)
    if not backup:
        error_count += 1
        return record_list, error_count
    if scan_result['owner'] == 'fail':
        if modify_file_owner(path, func_args):
            record_list.append({'info': 'Reinforce file {} owner success'.format(path), 'type': 'info'})
        else:
            error_count += 1
            record_list.append({'info': 'Reinforce file {} owner failed'.format(path), 'type': 'error'})
    if scan_result['priv'] == 'fail':
        if modify_file_privilege(path, func_args):
            record_list.append({'info': 'Reinforce file {} privilege success'.format(path), 'type': 'info'})
        else:
            error_count += 1
            record_list.append({'info': 'Reinforce file {} privilege failed'.format(path), 'type': 'error'})
    return record_list, error_count


"""
module: docker file host
func: 
    1. backup audit.rules
    2. rollback audit.rules
"""


def backup_audit_rules(func_args):
    error_count = 0
    path = func_args.get('path', None)
    target_backup = "{}.{}.`date +%Y-%m-%d`".format(path, backup_suffix)
    # 检查是否存在可能被覆盖的备份文件
    cmd = "cp {} {}".format(path, target_backup)
    result, output = CBBCommonUtils.cbb_run_cmd(func_args['ip'], cmd, username=func_args['username'],
                                                passwd=func_args['passwd'])
    if not result:
        # 检查是否备份成功
        check_cmd = "ls {}".format(target_backup)
        check_result, check_output = CBBCommonUtils.cbb_run_cmd(func_args['ip'], check_cmd,
                                                                username=func_args['username'],
                                                                passwd=func_args['passwd'])
        if not check_result and len(check_output) != 0:
            des = 'Backup Succeed.'
        else:
            error_count += 1
            des = "Backup check FAILED."
        return des, error_count
    else:
        error_count += 1
        des = "Backup check FAILED."
        return des, error_count


def rollback_audit_rules(func_args):
    error_count = 0
    path = func_args.get('path', None)
    real_path, filename = os.path.split(path)
    target_backup = "{}.{}*".format(path, backup_suffix)
    # 查找备份文件，此处的查找支持多备份文件的查找（结果会按时间顺序排序），并默认回滚最新一次备份的内容，可以连续多次回滚
    grep_cmd = "ls {}".format(target_backup)
    result, output = CBBCommonUtils.cbb_run_cmd(func_args['ip'], grep_cmd, username=func_args['username'],
                                                passwd=func_args['passwd'])
    if not result and len(output) != 0:
        target = os.path.join(real_path, output[len(output) - 1].replace("\n", ""))
        cmd = "/bin/cp -rf {} {}".format(target, path)
        # 回滚文件，
        result, output = CBBCommonUtils.cbb_run_cmd(func_args['ip'], cmd, username=func_args['username'],
                                                    passwd=func_args['passwd'])
        if not result:
            del_cmd = "rm -rf {}".format(target)
            # 删除对应的备份文件
            result, output = CBBCommonUtils.cbb_run_cmd(func_args['ip'], del_cmd, username=func_args['username'],
                                                        passwd=func_args['passwd'])
            if not result:
                del_des = "Backup file deleted. Target is {}.".format(target)
            else:
                error_count += 1
                del_des = "Backup file delete FAILED."
            return del_des, error_count
        else:
            error_count += 1
            des = "Rollback FAILED, please retry."
            return des, error_count
    else:
        error_count = -1
        des = "Backup file NOT FOUND."
        return des, error_count


"""
module: daemon file config
func: 
    1. validate daemon process is started
    2. validate args
"""


def validate_daemon_process_running(func_args):
    cmd = 'docker ps --quiet'
    result, output = CBBCommonUtils.cbb_run_cmd(func_args['ip'], cmd, username=func_args['username'],
                                                passwd=func_args['passwd'])
    if result:
        des = "CMD process ERROR. Error info: {}".format(result)
        return False, des
    if output and len(output):
        return True, output
    else:
        des = 'No images in docker, or the image in docker is not running.'
        return False, des


def scan_daemon_process_config(parse_config, func_args):
    is_valid, des = validate_daemon_process_running(func_args)
    if not is_valid:
        return [{'info': des, 'type': 'error'}], 1
    scan_list = []
    error_count = 0
    variable = func_args.get('variable', None)
    grep = func_args.get('grep', None)
    cmd = "docker ps --quiet --all | xargs docker inspect --format '{}'".format(variable)
    if grep:
        cmd = "docker ps --quiet | xargs docker inspect --format '{}' |grep -w {}".format(variable, grep)
    result, output = CBBCommonUtils.cbb_run_cmd(func_args['ip'], cmd, username=func_args['username'],
                                                passwd=func_args['passwd'])
    if result:
        error_count += 1
        scan_list.append({'info': "CMD process ERROR. Error info: {}".format(result), 'type': 'error'})
        return scan_list, error_count
    if not output or not len(output):
        if not grep:
            error_count += 1
            scan_list.append({'info': "The container in docker is not running.", 'type': 'warning'})
            return scan_list, error_count
        else:
            des = 'Safe. No containers were mapped in docker.sock'
            scan_list.append({'info': des, 'type': 'info'})
            scan_list.append({'info': 'Scan Complete.', 'type': 'info'})
            return scan_list, error_count
    for item in output:
        parse_info, parse_count = parse_config(item)
        scan_list.append(parse_info)
        error_count = error_count + parse_count
    scan_list.append({'info': 'Scan Complete.', 'type': 'info'})
    return scan_list, error_count


"""
module: daemon process config
func: 
    1. validate docker daemon process is started
    2. validate args
"""


def is_dockerd_exist(func_args):
    cmd = "ps -ef |grep dockerd |grep -v grep"
    result, output = CBBCommonUtils.cbb_run_cmd(func_args['ip'], cmd, username=func_args['username'],
                                                passwd=func_args['passwd'])
    if result:
        des = "CMD process ERROR. Error info: {}".format(result)
        return False, des
    if output and len(output) >= 1:
        return True, output
    else:
        des = 'Not Safe. Dockerd process is not existed.'
        return False, des


def scan_docker_process(func_args, parse_config):
    cmd = func_args.get('cmd', None)
    result, output = CBBCommonUtils.cbb_run_cmd(func_args['ip'], cmd, username=func_args['username'],
                                                passwd=func_args['passwd'])
    if result:
        des = "CMD process ERROR. Error info: {}".format(result)
        return des, 1
    des, error = parse_config(output)
    return des, error


"""
module: docker check main
func: check execute
"""


def check_func(scan, rollback, func_args):
    comm = func_args.get('comm', None)
    if comm == 1:
        des_list, error_scan = scan(func_args, flag=0)
        step_error = int(error_scan)
    elif comm == 2:
        des_list, error_reinforce = scan(func_args, flag=1)
        step_error = int(error_reinforce)
    elif comm == 3:
        des_list, error_rollback = rollback(func_args)
        step_error = int(error_rollback)
    else:
        return {"code": 3, "count": 0, "des": ['command must be 1/2/3']}
    if step_error == 0:
        code = 0
    elif step_error <= -1:
        code = 2
    else:
        code = 1
    return {"code": code, "count": step_error, "des": des_list}
